import cv2
import time
from rknnpool import rknnPoolExecutor
# 图像处理函数，实际应用过程中需要自行修改
from func import myFunc
import argparse


def handel_video(opt):
    cap = cv2.VideoCapture('./720p60hz.mp4')
    frame_name = 1  # 使用数字命名将要被保存的帧
    fps = cap.get(cv2.CAP_PROP_FPS)  # 获取视频的帧率
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  # 获取视频文件的帧宽度
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 获取视频文件的帧高度

    # 创建VideoWriter对象以保存提取的帧为新的视频文件
    output = cv2.VideoWriter('out.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
    # cap = cv2.VideoCapture(0)
    modelPath = "./rknnModel/yolov5s_rk3588.rknn"
    # 线程数, 增大可提高帧率 推荐为6
    TPEs = opt.tpes
    # 初始化rknn池
    pool = rknnPoolExecutor(
        rknnModel=modelPath,
        TPEs=TPEs,
        func=myFunc)

    # 初始化异步所需要的帧
    if (cap.isOpened()):
        for i in range(TPEs + 1):
            ret, frame = cap.read()
            if not ret:
                cap.release()
                del pool
                exit(-1)
            pool.put(frame)

    frames, loopTime, initTime = 0, time.time(), time.time()
    while (cap.isOpened()):
        frames += 1
        ret, frame = cap.read()
        if not ret:
            break
        pool.put(frame)
        frame, flag = pool.get()
        if flag == False:
            break
        # cv2.imshow('test', frame)
        # cv2.imwrite("./img/" + str(frame_name) + ".jpg", frame)  # 保存视频帧
        # frame_name = frame_name + 1  # 用于命名将要被保存的帧的数字执行自加操作
        output.write(frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
        if frames % 30 == 0:
            print("30帧平均帧率:\t", 30 / (time.time() - loopTime), "帧")
            loopTime = time.time()

    print("总平均帧率\t", frames / (time.time() - initTime))
    # 释放资源
    release(cap, output, pool)


def release(cap, output, pool):
    cap.release()
    output.release()
    # cv2.destroyAllWindows()
    pool.release()


def parse_opt(known=False):
    parser = argparse.ArgumentParser()
    parser.add_argument('--tpes', type=int, default=1, help='tpes')
    opt = parser.parse_known_args()[0] if known else parser.parse_args()
    return opt


if __name__ == '__main__':
    opt = parse_opt()
    handel_video(opt)
